import asyncio
import websockets
import json
import time
import argparse
import statistics
from typing import Dict, List, Any

class WebSocketTest:
    def __init__(self, host='127.0.0.1', port=5678):
        self.uri = f'ws://{host}:{port}'
        self.connected_clients: Dict[int, websockets.WebSocketClientProtocol] = {}
        self.locks: Dict[int, asyncio.Lock] = {}
        self.active_tasks: List[asyncio.Task] = []

    async def connect_client(self, client_id: int) -> websockets.WebSocketClientProtocol:
        """建立单个连接并保持心跳"""
        try:
            ws = await websockets.connect(self.uri)
            self.connected_clients[client_id] = ws
            self.locks[client_id] = asyncio.Lock()

            # 发送登录消息
            async with self.locks[client_id]:
                login_msg = {
                    "c": "user",
                    "m": "login",
                    "data": {"uid": client_id}
                }
                await ws.send(json.dumps(login_msg))
                response = await ws.recv()
                response_data = json.loads(response)
                if response_data.get('code', -1) != 0:
                    raise Exception(f"登录失败: {response_data}")

            print(f"Client {client_id} connected and logged in")

            # 启动心跳协程并保存任务引用
            heartbeat_task = asyncio.create_task(self.keep_alive(ws, client_id))
            self.active_tasks.append(heartbeat_task)
            return ws

        except Exception as e:
            print(f"Client {client_id} connection error: {str(e)}")
            if client_id in self.connected_clients:
                del self.connected_clients[client_id]
            if client_id in self.locks:
                del self.locks[client_id]
            return None

    async def keep_alive(self, ws: websockets.WebSocketClientProtocol, client_id: int) -> None:
        """心跳保持协程"""
        try:
            while True:
                try:
                    async with self.locks[client_id]:
                        ping_msg = {
                            "c": "socket",
                            "m": "ping",
                            "data": {"uid": client_id}
                        }
                        await ws.send(json.dumps(ping_msg))
                        # response = await asyncio.wait_for(ws.recv(), timeout=5.0)
                        # response_data = json.loads(response)
                        # if response_data.get('code', -1) != 0:
                        #     print(f"Client {client_id} heartbeat failed: {response_data}")
                        #     continue

                    # 心跳间隔改为10秒，避免太频繁
                    await asyncio.sleep(10)

                except asyncio.TimeoutError:
                    print(f"Client {client_id} heartbeat timeout, retrying...")
                    await asyncio.sleep(1)
                    continue
                except websockets.exceptions.ConnectionClosed:
                    print(f"Client {client_id} connection closed")
                    break
                except Exception as e:
                    print(f"Client {client_id} heartbeat error: {str(e)}")
                    await asyncio.sleep(1)
                    continue

        except asyncio.CancelledError:
            print(f"Client {client_id} heartbeat stopped")
        except Exception as e:
            print(f"Client {client_id} heartbeat error: {str(e)}")
        finally:
            await self.cleanup_client(client_id)

    async def cleanup_client(self, client_id: int) -> None:
        """清理单个客户端连接"""
        if client_id in self.connected_clients:
            try:
                await self.connected_clients[client_id].close()
            except:
                pass
            del self.connected_clients[client_id]
        if client_id in self.locks:
            del self.locks[client_id]

    async def cpu_test(self, ws: websockets.WebSocketClientProtocol, client_id: int, count: int) -> Dict[str, Any]:
        """CPU密集任务测试"""
        success = 0
        failed = 0
        response_times = []

        for i in range(count):
            try:
                async with self.locks[client_id]:
                    start_time = time.time()
                    message = {
                        "c": "user",
                        "m": "cpu",
                        "data": {"uid": client_id}
                    }
                    await ws.send(json.dumps(message))
                    response = await asyncio.wait_for(ws.recv(), timeout=30.0)
                    end_time = time.time()

                    response_data = json.loads(response)
                    print(f"Client {client_id} request {i} response: {response_data}")
                    if response_data.get('code', -1) == 0:
                        success += 1
                        response_time = (end_time - start_time) * 1000
                        response_times.append(response_time)
                        print(f"Client {client_id} request {i} success, response time: {response_time:.2f} ms")
                    else:
                        failed += 1
                        print(f"Client {client_id} request {i} failed: {response_data}")

                await asyncio.sleep(0.01)

            except Exception as e:
                print(f"Client {client_id} request {i} error: {str(e)}")
                failed += 1

        return {
            "success": success,
            "failed": failed,
            "response_times": response_times
        }

    async def run_performance_test(self, num_clients: int, requests_per_client: int) -> None:
        """执行性能测试"""
        print(f"\n开始连接测试: {num_clients} 个客户端")
        start_time = time.time()

        # 建立连接
        connect_tasks = []
        for i in range(num_clients):
            task = asyncio.create_task(self.connect_client(1000 + i))
            connect_tasks.append(task)
            await asyncio.sleep(0.1)

        connections = await asyncio.gather(*connect_tasks)
        active_connections = [c for c in connections if c is not None]

        if not active_connections:
            print("没有成功建立的连接")
            return

        connect_duration = time.time() - start_time
        print(f"\n连接测试完成:")
        print(f"成功连接数: {len(active_connections)}")
        print(f"连接耗时: {connect_duration:.2f} 秒")
        print(f"连接速率: {len(active_connections)/connect_duration:.2f} 连接/秒")

        await asyncio.sleep(2)

        # 执行CPU测试
        print(f"\n开始CPU测试: 每个客户端发送 {requests_per_client} 个请求")
        test_start_time = time.time()

        test_tasks = []
        for client_id, ws in self.connected_clients.items():
            task = asyncio.create_task(self.cpu_test(ws, client_id, requests_per_client))
            test_tasks.append(task)

        try:
            results = await asyncio.gather(*test_tasks)
        except Exception as e:
            print(f"测试执行错误: {str(e)}")
            return

        # 统计结果
        total_success = sum(r["success"] for r in results)
        total_failed = sum(r["failed"] for r in results)
        all_response_times = []
        for r in results:
            all_response_times.extend(r["response_times"])

        test_duration = time.time() - test_start_time

        # 输出详细统计
        print("\n性能测试结果:")
        print(f"总请求数: {total_success + total_failed}")
        print(f"成功请求数: {total_success}")
        print(f"失败请求数: {total_failed}")
        print(f"成功率: {(total_success/(total_success + total_failed))*100:.2f}%")
        print(f"总耗时: {test_duration:.2f} 秒")
        print(f"吞吐量: {total_success/test_duration:.2f} 请求/秒")

        if all_response_times:
            print("\n响应时间统计:")
            print(f"平均响应时间: {statistics.mean(all_response_times):.2f} ms")
            print(f"最小响应时间: {min(all_response_times):.2f} ms")
            print(f"最大响应时间: {max(all_response_times):.2f} ms")
            print(f"响应时间中位数: {statistics.median(all_response_times):.2f} ms")
            print(f"响应时间90百分位: {sorted(all_response_times)[int(len(all_response_times)*0.9)]:.2f} ms")
            print(f"响应时间95百分位: {sorted(all_response_times)[int(len(all_response_times)*0.95)]:.2f} ms")
            print(f"响应时间99百分位: {sorted(all_response_times)[int(len(all_response_times)*0.99)]:.2f} ms")

    async def cleanup(self) -> None:
        """清理所有连接和任务"""
        # 取消所有活动任务
        for task in self.active_tasks:
            task.cancel()

        # 等待任务完成
        if self.active_tasks:
            await asyncio.gather(*self.active_tasks, return_exceptions=True)

        # 关闭所有连接
        for client_id in list(self.connected_clients.keys()):
            await self.cleanup_client(client_id)

async def main():
    parser = argparse.ArgumentParser(description='WebSocket Performance Test')
    parser.add_argument('--host', default='127.0.0.1', help='WebSocket server host')
    parser.add_argument('--port', type=int, default=5678, help='WebSocket server port')
    parser.add_argument('--clients', type=int, default=10, help='Number of clients')
    parser.add_argument('--requests', type=int, default=100, help='Requests per client')

    args = parser.parse_args()
    test = WebSocketTest(args.host, args.port)

    try:
        await test.run_performance_test(args.clients, args.requests)
        print("\n按 Ctrl+C 停止测试...")
        while True:
            await asyncio.sleep(1)
    except KeyboardInterrupt:
        print("\n正在清理连接...")
        await test.cleanup()
        print("测试完成")
    except Exception as e:
        print(f"\n测试过程中出现错误: {str(e)}")
        await test.cleanup()

if __name__ == "__main__":
    asyncio.run(main())